package engine

import (
	"context"
	"fmt"
	"sync"
	"time"

	"gitee.com/zfd81/dbrs/sql/annotation"

	"golang.org/x/sync/errgroup"

	"github.com/opentracing/opentracing-go"
)

var ctxNowFunc = time.Now
var ctxNowFuncMutex = &sync.Mutex{}

func RunWithNowFunc(nowFunc func() time.Time, fn func() error) error {
	ctxNowFuncMutex.Lock()
	defer ctxNowFuncMutex.Unlock()

	initialNow := ctxNowFunc
	ctxNowFunc = nowFunc
	defer func() {
		ctxNowFunc = initialNow
	}()

	return fn()
}

type Context struct {
	context.Context
	Session
	pid         uint64
	query       string
	sql         string
	queryTime   time.Time
	tracer      opentracing.Tracer
	rootSpan    opentracing.Span
	annotations annotation.Annotations
}

type ContextOption func(*Context)

func NewContext(
	ctx context.Context,
	opts ...ContextOption,
) *Context {
	c := &Context{
		Context:   ctx,
		Session:   NewBaseSession(),
		queryTime: ctxNowFunc(),
		tracer:    opentracing.NoopTracer{},
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

func NewEmptyContext() *Context { return NewContext(context.TODO()) }

// Applys the options given to the context. Mostly for tests, not safe for use after construction of the context.
func (c *Context) ApplyOpts(opts ...ContextOption) {
	for _, opt := range opts {
		opt(c)
	}
}

// Pid returns the process id associated with this context.
func (c *Context) Pid() uint64 { return c.pid }

// Query returns the query string associated with this context.
func (c *Context) Query() string { return c.query }

// Sql returns the sql string associated with this context.
func (c *Context) Sql() string { return c.sql }

// QueryTime returns the time.Time when the context associated with this query was created
func (c *Context) QueryTime() time.Time {
	return c.queryTime
}

func (c *Context) GetAnnotation(atype annotation.Type) []annotation.Annotation {
	return c.annotations.Get(atype)
}

func (c *Context) GetAllAnnotations() annotation.Annotations {
	return c.annotations
}

// Span creates a new tracing span with the given context.
// It will return the span and a new context that should be passed to all
// children of this span.
func (c *Context) Span(
	opName string,
	opts ...opentracing.StartSpanOption,
) (opentracing.Span, *Context) {
	parentSpan := opentracing.SpanFromContext(c.Context)
	if parentSpan != nil {
		opts = append(opts, opentracing.ChildOf(parentSpan.Context()))
	}
	span := c.tracer.StartSpan(opName, opts...)
	ctx := opentracing.ContextWithSpan(c.Context, span)

	return span, c.WithContext(ctx)
}

// NewSubContext creates a new sub-context with the current context as parent. Returns the resulting context.CancelFunc
// as well as the new *sql.Context, which be used to cancel the new context before the parent is finished.
func (c *Context) NewSubContext() (*Context, context.CancelFunc) {
	ctx, cancelFunc := context.WithCancel(c.Context)

	return c.WithContext(ctx), cancelFunc
}

// WithContext returns a new context with the given underlying context.
func (c *Context) WithContext(ctx context.Context) *Context {
	nc := *c
	nc.Context = ctx
	return &nc
}

// RootSpan returns the root span, if any.
func (c *Context) RootSpan() opentracing.Span {
	return c.rootSpan
}

// Error adds an error as warning to the session.
func (c *Context) Error(code int, msg string, args ...interface{}) {
	c.Session.Warn(&Warning{
		Level:   "Error",
		Code:    code,
		Message: fmt.Sprintf(msg, args...),
	})
}

// Warn adds a warning to the session.
func (c *Context) Warn(code int, msg string, args ...interface{}) {
	c.Session.Warn(&Warning{
		Level:   "Warning",
		Code:    code,
		Message: fmt.Sprintf(msg, args...),
	})
}

func (c *Context) NewErrgroup() (*errgroup.Group, *Context) {
	eg, egCtx := errgroup.WithContext(c.Context)
	return eg, c.WithContext(egCtx)
}

// WithSession adds the given session to the context.
func WithSession(s Session) ContextOption {
	return func(ctx *Context) {
		ctx.Session = s
	}
}

// WithTracer adds the given tracer to the context.
func WithTracer(t opentracing.Tracer) ContextOption {
	return func(ctx *Context) {
		ctx.tracer = t
	}
}

// WithPid adds the given pid to the context.
func WithPid(pid uint64) ContextOption {
	return func(ctx *Context) {
		ctx.pid = pid
	}
}

// WithQuery adds the given query to the context.
func WithQuery(q string) ContextOption {
	return func(ctx *Context) {
		ctx.sql, ctx.annotations = annotation.Parser(q)
		ctx.query = q
	}
}

// WithRootSpan sets the root span of the context.
func WithRootSpan(s opentracing.Span) ContextOption {
	return func(ctx *Context) {
		ctx.rootSpan = s
	}
}
